KEY_PREFIX = 'tf2x_cifar10_'
UINT8_SUFFIX = '_uint8'
FLT32_SUFFIX = '_flt32'

if '__main__' == __name__:
    import datetime
    t1 = datetime.datetime.now()
    import tensorflow as tf
    from python_ai.category.redis.conn.local_redis import r as redis
    from python_ai.common.read_data.redis_numpy import toRedisNdLg, fromRedisNdLg
    import numpy as np
    from python_ai.common.xcommon import *
    import sys
    import pickle

    sep('Names of classes ...')
    BASE_DIR, FILE_NAME = os.path.split(__file__)
    path = '../../../large_data/DL1/cifar-10-batches-py/batches.meta'
    FILE_PATH = os.path.join(BASE_DIR, path)
    with open(FILE_PATH, 'br') as f:
        xdict = pickle.load(f)
    print(xdict)
    names_list = np.array(xdict['label_names'])
    toRedisNdLg(redis, names_list, KEY_PREFIX + 'names')
    print('Names OK')

    print('Loading data ...')
    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
    print_numpy_ndarray_info_simple(x_train, 'x_train')
    print_numpy_ndarray_info_simple(x_test, 'x_test')
    print_numpy_ndarray_info_simple(y_train, 'y_train')
    print_numpy_ndarray_info_simple(y_test, 'y_test')
    print('Loaded.')

    print('Shuffling ...')
    # shuffle training data
    n_train = len(y_train)
    rnd_idx = np.random.permutation(n_train)
    x_train = x_train[rnd_idx]
    y_train = y_train[rnd_idx]
    # shuffle testing data
    n_test = len(y_test)
    rnd_idx = np.random.permutation(n_test)
    x_test = x_test[rnd_idx]
    y_test = y_test[rnd_idx]
    print('Done')

    print('Split test into test and val')
    n_val = n_test // 2
    x_test_ = x_test
    y_test_ = y_test
    x_test = x_test_[:n_val]
    y_test = y_test_[:n_val]
    x_val = x_test_[n_val:]
    y_val = y_test_[n_val:]
    print_numpy_ndarray_info_simple(x_train, 'x_train')
    print_numpy_ndarray_info_simple(x_test, 'x_test')
    print_numpy_ndarray_info_simple(x_test, 'x_val')
    print_numpy_ndarray_info_simple(y_train, 'y_train')
    print_numpy_ndarray_info_simple(y_test, 'y_test')
    print_numpy_ndarray_info_simple(y_test, 'y_val')
    print('Done')

    print('LUT uint8 => flt32')
    lut = np.arange(256, dtype=np.float32) / 255.
    x_train_flt32 = lut[x_train]
    x_test_flt32 = lut[x_test]
    x_val_flt32 = lut[x_val]
    t2 = datetime.datetime.now()
    print(f'Loaded and processed. ({t2 - t1})')
    print_numpy_ndarray_info_simple(x_train, 'x_train_uint8')
    print_numpy_ndarray_info_simple(x_train_flt32, 'x_train_flt32')
    print_numpy_ndarray_info_simple(x_test, 'x_test_uint8')
    print_numpy_ndarray_info_simple(x_test_flt32, 'x_test_flt32')
    print_numpy_ndarray_info_simple(x_val, 'x_val_uint8')
    print_numpy_ndarray_info_simple(x_val_flt32, 'x_val_flt32')
    print('OK')

    print('Storing to redis ori data ...')
    toRedisNdLg(redis, x_train, KEY_PREFIX + 'x_train' + UINT8_SUFFIX)
    toRedisNdLg(redis, y_train, KEY_PREFIX + 'y_train')
    toRedisNdLg(redis, x_test, KEY_PREFIX + 'x_test' + UINT8_SUFFIX)
    toRedisNdLg(redis, y_test, KEY_PREFIX + 'y_test')
    toRedisNdLg(redis, x_val, KEY_PREFIX + 'x_val' + UINT8_SUFFIX)
    toRedisNdLg(redis, y_val, KEY_PREFIX + 'y_val')
    print('OK')

    print('Storing to redis flt32 data ...')
    toRedisNdLg(redis, x_train_flt32, KEY_PREFIX + 'x_train' + FLT32_SUFFIX)
    toRedisNdLg(redis, x_test_flt32, KEY_PREFIX + 'x_test' + FLT32_SUFFIX)
    toRedisNdLg(redis, x_val_flt32, KEY_PREFIX + 'x_val' + FLT32_SUFFIX)
    print('OK')

    print(f'Over, tf2x keras cifar10 is put into redis {redis}')
